import os
import json
from torchvision.ops import box_iou
import torch
origin_json_path = "/home/jxy/datasets/gen-vlkt/hico_20160224_det/annotations/trainval_hico.json"
# generate_hico_dir = "./test-add_seg-max_adapt-black-dynamic_threshold"
generate_hico_dir = '/home/jxy/program/hoi2/600HOI-600HOI_no_object-add_seg_only_instance-max_adapt-dynamic_threshold_add_weight_23-black_add_correlation_hoi/generate_hoi/generate_hoi'

with open(origin_json_path,'r',encoding='utf8') as f:
    json_data = json.load(f)
    # json_data = eval(json_data)
    print(type(json_data))
    print(len(json_data))
    # exit()
    count_temp = 0
    origin_total_hoi_count = 0
    origin_total_boxes_count = 0
    generate_total_hoi_count = 0
    generate_total_boxes_count = 0
    correct_match_hoi_count = 0
    correct_match_boxes_count = 0
    all_include_hoi_count = 0
    ge_half_include_hoi_count = 0
    l_half_include_hoi_count = 0
    zero_include_hoi_count = 0
    total_image_count = 0
    count_ = 1
    for hoi in json_data:
        if count_ > 1000000:
            break
        else:
            count_ += 1

        # if os.path.basename(hoi["file_name"]) != "HICO_train2015_00000009.jpg":
        #     continue
        print(hoi)
        total_image_count += 1
        file_name = hoi["file_name"]
        txt_name = file_name.split(".")[0] + ".txt"
        print(txt_name)
        with open(os.path.join(generate_hico_dir, txt_name), "r") as f:
            generate_hoi_dict = eval(f.read())
        origin_boxes = hoi["annotations"]
        origin_hoi = hoi["hoi_annotation"]
        generate_boxes = generate_hoi_dict["annotations"]
        generate_hoi = generate_hoi_dict["hoi_annotation"]

        print(origin_boxes)
        print(origin_hoi)
        print(generate_boxes)
        print(generate_hoi)
        origin_total_hoi_count += len(origin_hoi)
        generate_total_hoi_count += len(generate_hoi)

        origin_total_boxes_count += len(origin_boxes)
        generate_total_boxes_count += len(generate_boxes)

        origin2generate_index_dict = {}
        for i in range(len(origin_boxes)):
            box_i = origin_boxes[i]["bbox"]
            max_iou = -1
            max_j = -1
            for j in range(len(generate_boxes)):
                box_j = generate_boxes[j]["bbox"]
                iou_ = box_iou(torch.tensor(box_i).unsqueeze(0), torch.tensor(box_j).unsqueeze(0))
                print(iou_)
                if max_iou < iou_:
                    max_iou = iou_
                    max_j = j
            if max_iou > 0.5 and origin_boxes[i]["category_id"] == generate_boxes[max_j]["category_id"]:
                origin2generate_index_dict[i] = max_j
                correct_match_boxes_count += 1
            else:
                origin2generate_index_dict[i] = -1
        print(origin2generate_index_dict)
        print(correct_match_boxes_count)

        current_find_count = 0
        for i in range(len(origin_hoi)):
            hoi_i = origin_hoi[i]
            hoi_i_subject_id = hoi_i["subject_id"]
            hoi_i_object_id = hoi_i["object_id"]
            hoi_i_category_id = hoi_i["category_id"]
            hoi_i_hoi_category_id = hoi_i["hoi_category_id"]
            if origin2generate_index_dict[hoi_i_subject_id] == -1 or \
               origin2generate_index_dict[hoi_i_object_id] == -1:
                continue
            flag = False
            for j in range(len(generate_hoi)):
                hoi_j = generate_hoi[j]
                hoi_j_subject_id = hoi_j["subject_id"]
                hoi_j_object_id = hoi_j["object_id"]
                hoi_j_category_id = hoi_j["category_id"]
                hoi_j_hoi_category_id = hoi_j["hoi_category_id"]
                if origin2generate_index_dict[hoi_i_subject_id] == hoi_j_subject_id and \
                   origin2generate_index_dict[hoi_i_object_id] == hoi_j_object_id and \
                   hoi_i_category_id == hoi_j_category_id and \
                   hoi_i_hoi_category_id == hoi_j_hoi_category_id:
                    correct_match_hoi_count += 1
                    current_find_count += 1
                    flag = True
                    break
        # if not flag:
        # 统计全包括的，大于一半的，小于一半的关系的个数。
        if current_find_count == 0:
            zero_include_hoi_count += 1
        elif current_find_count == len(origin_hoi):
            all_include_hoi_count += 1
        elif current_find_count / len(origin_hoi) >= 0.5:
            ge_half_include_hoi_count += 1
        else:
            l_half_include_hoi_count += 1
    print("=========================================================")
    # print("ground_truth所有关系数:", origin_total_hoi_count)
    print("共%d张图像" % total_image_count)
    print("=========================================================")
    print("ground_truth所有关系数:", origin_total_hoi_count)
    print("generate_ground_truth所有关系数:", generate_total_hoi_count)
    print("生成正确关系的数量:", correct_match_hoi_count, "占原比:", round(correct_match_hoi_count / origin_total_hoi_count, 2), "占生成比:", round(correct_match_hoi_count / generate_total_hoi_count, 2))
    print("其中关系全部包括的:", all_include_hoi_count, "占比:", round(all_include_hoi_count / total_image_count, 2))
    print("其中关系包括大于等于一半的:", ge_half_include_hoi_count, "占比:", round(ge_half_include_hoi_count / total_image_count, 2))
    print("其中关系包括小于一半的:", l_half_include_hoi_count, "占比:", round(l_half_include_hoi_count / total_image_count, 2))
    print("一个都不对的:", zero_include_hoi_count, "占比:", round(zero_include_hoi_count / total_image_count, 2))
    print("=========================================================")
    # print("=========================================================")
    print("ground_truth所有检测框数:", origin_total_boxes_count)
    print("generate_ground_truth所有检测框数:", generate_total_boxes_count)
    print("正确检测框个数:", correct_match_boxes_count, "占原比:", round(correct_match_boxes_count / origin_total_boxes_count, 2), "占生成比:", round(correct_match_boxes_count / generate_total_boxes_count, 2))
    print("=========================================================")
    exit()
                    
    exit()

